#include "sql_writer_pq.h"
#include <libpq-fe.h>
#include <unordered_map>
#include <vector>
#include <sstream>
#include <iomanip>
#include <chrono>
#include <cstring>
#include <iostream>

using namespace kwdb::ingest::v1;

static const char* env(const char* k, const char* defv){ const char* v=getenv(k); return v&&*v?v:defv; }

// Thread-safe connection creation (each request gets its own connection)
static PGconn* create_conn(std::string* err){
  std::ostringstream ci;
  ci<<"host="<<env("KWDB_HOST","127.0.0.1")
    <<" port="<<env("KWDB_SQL_PORT","26257")
    <<" dbname="<<env("KWDB_DB","kwdb")
    <<" user="<<env("KWDB_USER","root")
    <<" sslmode=disable";
  const char* pwd=getenv("KWDB_PWD"); if (pwd && *pwd) ci<<" password="<<pwd;

  PGconn* conn = PQconnectdb(ci.str().c_str());
  if (PQstatus(conn)!=CONNECTION_OK){
    if(err)*err=PQerrorMessage(conn);
    std::cerr << "[sql_writer_pq] Connection failed: " << (err?*err:"unknown") << std::endl;
    return nullptr;
  }

  return conn;
}

static std::string quote_ident(PGconn* c, const std::string& s){
  char* q = PQescapeIdentifier(c, s.c_str(), s.size());
  std::string out = q?q:"\"\""; if(q) PQfreemem(q); return out;
}

static std::string quote_literal(PGconn* c, const std::string& s){
  char* q = PQescapeLiteral(c, s.c_str(), s.size());
  std::string out = q?q:"''"; if(q) PQfreemem(q); return out;
}

static std::string ts_ns_to_sql_literal(int64_t ns, PGconn* c){
  time_t sec = ns/1000000000LL; long nsec = ns%1000000000LL;
  std::tm tm; gmtime_r(&sec, &tm);
  std::ostringstream oss;
  oss<<std::put_time(&tm, "%Y-%m-%d %H:%M:%S")<<'.'
     <<std::setw(9)<<std::setfill('0')<<nsec;
  return quote_literal(c, oss.str());
}

static std::unordered_map<std::string, std::vector<const Point*>> group_by_meas(const WriteRequest& req){
  std::unordered_map<std::string, std::vector<const Point*>> g;
  for (const auto& p: req.points()) g[p.measurement()].push_back(&p);
  return g;
}

static bool exec_ok(PGconn* c, const char* sql, std::string* err){
  PGresult* r = PQexec(c, sql);
  auto st = PQresultStatus(r);
  if (st!=PGRES_COMMAND_OK && st!=PGRES_TUPLES_OK){
    if(err)*err = PQerrorMessage(c);
    PQclear(r); return false;
  }
  PQclear(r); return true;
}

long long write_points_sql(const WriteRequest& req, std::string* err){
  auto start = std::chrono::steady_clock::now();
  int total_points = req.points_size();

  std::cout << "[sql_writer_pq] recv " << total_points << " points" << std::endl;

  // Create new connection for this request (thread-safe)
  PGconn* c = create_conn(err);
  if (!c) return -1;

  auto groups = group_by_meas(req);
  long long total = 0;

  for (auto& kv: groups){
    const std::string& meas = kv.first;
    const auto& vec = kv.second;
    if (vec.empty()) continue;

    // 1) Collect union of all columns (tags/fields)
    std::vector<std::string> tag_cols, field_cols;
    {
      std::unordered_map<std::string,int> seen;
      for (auto* p: vec){
        for (const auto& t: p->tags())   if(!seen.count(t.key()))   { seen[t.key()]=1; tag_cols.push_back(t.key()); }
        for (const auto& f: p->fields()) if(!seen.count("\tF"+f.key())){ seen["\tF"+f.key()]=1; field_cols.push_back(f.key()); }
      }
    }

    // 2) Build INSERT with multiple values
    std::ostringstream sql;
    sql<<"INSERT INTO "<<quote_ident(c, meas)<<"("<<quote_ident(c,"ts");
    for (auto& col: tag_cols)   sql<<","<<quote_ident(c,col);
    for (auto& col: field_cols) sql<<","<<quote_ident(c,col);
    sql<<") VALUES ";

    bool first_row = true;
    const size_t max_vals_per_stmt = 800; // Prevent SQL from being too long, adjustable
    size_t row_count_in_stmt = 0;

    auto flush_stmt = [&]()->bool{
      if (row_count_in_stmt==0) return true;
      std::string s = sql.str(); sql.str(""); sql.clear();
      if (!exec_ok(c, s.c_str(), err)) return false;
      row_count_in_stmt = 0; return true;
    };

    if (!exec_ok(c, "BEGIN", err)) return -1;

    for (auto* p: vec){
      if (!first_row) sql<<","; first_row=false;
      row_count_in_stmt++;

      // Current row values
      std::unordered_map<std::string,std::string> tags;
      for (const auto& t: p->tags()) tags[t.key()] = t.value();
      std::unordered_map<std::string,double> fields;
      for (const auto& f: p->fields()) fields[f.key()] = f.f64();

      sql<<"("<<ts_ns_to_sql_literal(p->timestamp_unix_ns(), c);
      for (auto& col: tag_cols){
        auto it=tags.find(col);
        if (it==tags.end()) sql<<",NULL";
        else sql<<","<<quote_literal(c,it->second);
      }
      for (auto& col: field_cols){
        auto it=fields.find(col);
        if (it==fields.end()) sql<<",NULL";
        else                  sql<<","<<it->second;
      }
      sql<<")";

      // Execute batch when threshold reached, rewrite header
      if (row_count_in_stmt >= max_vals_per_stmt){
        if (!flush_stmt()) { exec_ok(c,"ROLLBACK",nullptr); return -1; }
        // Rewrite INSERT header
        sql<<"INSERT INTO "<<quote_ident(c, meas)<<"("<<quote_ident(c,"ts");
        for (auto& col: tag_cols)   sql<<","<<quote_ident(c,col);
        for (auto& col: field_cols) sql<<","<<quote_ident(c,col);
        sql<<") VALUES ";
        first_row=true;
      }
    }

    if (!flush_stmt()) { exec_ok(c,"ROLLBACK",nullptr); return -1; }
    if (!exec_ok(c, "COMMIT", err)) return -1;

    total += (long long)vec.size();
  }

  auto end = std::chrono::steady_clock::now();
  auto duration_ms = std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count();

  std::cout << "[sql_writer_pq] written=" << total << " points, elapsed=" << duration_ms << "ms" << std::endl;

  // Close connection
  PQfinish(c);

  return total;
}
